package comp9417;

import java.util.*;

/**
 * Created by huy on 20/05/14.
 * This is called naive because it linearly evaluates all points.
 * A better way would be to use a kd-trees.
 */
public class NaiveNearestNeighbour {

    private final Dataset dataset;
    private final Set<Integer> indices;

    public NaiveNearestNeighbour(Dataset dataset, String predictionAttribute) {
        this.dataset = dataset;
        this.indices = getIncludeAttributeIndices(predictionAttribute);
    }

    public LinkedHashMap<Double, DataPoint> nearestWithDistances(int k, DataPoint p) {
        TreeMap<Double, DataPoint> distances = calculateDistances(p, dataset, indices);
        final LinkedHashMap<Double, DataPoint> nearest = new LinkedHashMap<Double, DataPoint>(k);
        getkNearest(k, distances, new Collector() {
            @Override
            public void collect(Map.Entry<Double, DataPoint> entry) {
                nearest.put(entry.getKey(), entry.getValue());
            }
        });
        return nearest;
    }

    public List<DataPoint> nearest(int k, DataPoint p) {
        final List<DataPoint> nearest = new ArrayList<DataPoint>(k);
        getkNearest(k, calculateDistances(p, dataset, indices), new Collector() {
            @Override
            public void collect(Map.Entry<Double, DataPoint> entry) {
                nearest.add(entry.getValue());
            }
        });
        return nearest;
    }

    private void getkNearest(int k, Map<Double, DataPoint> distances, Collector collector) {
        int count = k;
        for (Map.Entry<Double, DataPoint> entry : distances.entrySet()) {
            collector.collect(entry);
            count--;
            if (count == 0) break;
        }
    }

    private TreeMap<Double, DataPoint> calculateDistances(DataPoint p, Dataset dataset, Set<Integer> indices) {
        TreeMap<Double, DataPoint> distances = new TreeMap<Double, DataPoint>();
        for (DataPoint dataPoint : dataset.getDatapoints()) {
            distances.put(euclideanDistanceSquare(p, dataPoint, indices), dataPoint);
        }
        return distances;
    }

    private Set<Integer> getIncludeAttributeIndices(String predictionAttribute) {
        if (predictionAttribute == null) {
            throw new IllegalArgumentException("Invalid Prediction Attribute");
        }
        Set<Integer> indices = new HashSet<Integer>();
        for (int index = 0; index < dataset.getAttributes().length; index++) {
            Attribute attribute = dataset.getAttributes()[index];
            if (!predictionAttribute.equalsIgnoreCase(attribute.name) && !attribute.prune) {
                indices.add(index);
            }
        }
        if (indices.size() == dataset.getAttributes().length) {
            throw new RuntimeException("Prediction Attribute does not exist in dataset.");
        }
        return indices;
    }

    private double euclideanDistanceSquare(DataPoint p1, DataPoint p2, Set<Integer> indices) {
        if (p1.getValues().length != p2.getValues().length) {
            throw new IllegalArgumentException("DataPoint do not have the same dimensions");
        }
        double runningTotal = 0.0;
        for (int index : indices) {
            Object v1 = p1.getValues()[index];
            Object v2 = p2.getValues()[index];
            runningTotal += Math.pow(difference(v1, v2, index), 2.0);
        }
        return runningTotal;
    }

    private double difference(Object v1, Object v2, int index) {
        return dataset.getAttributes()[index].difference(v1, v2);
    }

    private interface Collector {
        public void collect(Map.Entry<Double, DataPoint> entry);
    }

}
